//
// Created by bobin on 17-11-15.
//

#include "Optimize.h"

Optimize::Optimize(const std::vector<Eigen::Vector3d> &_x3dPoints,
                   const std::vector<Eigen::Vector2d> &_x2dPoints, const Eigen::Matrix3d &_K,
                   const std::vector<int> &_pointCate,
                   const std::vector<int> &_pointLineIndex,
                   const std::vector<int> &_pointPlaneIndex,
                   const std::vector<g2o::Vector6D> &_lines,
                   const std::vector<Eigen::Vector4d> &_planes) {
    m3dPoints = _x3dPoints;
    m2dPoints = _x2dPoints;
    K = _K;
    mPointCate = _pointCate;
    mLines = _lines;
    mPlanes = _planes;
    mPointLineIndex = _pointLineIndex;
    mPointPlaneIndex = _pointPlaneIndex;
}

void Optimize::PnPL_plucker(const std::vector<cv::Point3f> &pts3d, const std::vector<cv::Point2f> &pts2d,
                            const std::vector<cv::Vec6f> &lns3d, const std::vector<cv::Vec4f> &lns2d,
                            const cv::Mat &K, cv::Mat &R, cv::Mat &t) {
    int npts = pts3d.size();
    int nlns = lns3d.size();

    // init g2o

    g2o::BlockSolver_6_6::LinearSolverType *linearSolver;
    linearSolver = new g2o::LinearSolverCSparse<g2o::BlockSolver_6_6::PoseMatrixType>();
    g2o::BlockSolver_6_6 *solver_ptr = new g2o::BlockSolver_6_6(linearSolver);
    g2o::OptimizationAlgorithmLevenberg *solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    g2o::SparseOptimizer optimizer;
    // optimizer.setVerbose(true);
    optimizer.setAlgorithm(solver);

    // add vertex for pose
    int vertex_id = 0;
    Eigen::Vector3d trans(0.1, 0.1, 0.1);
    Eigen::Quaterniond q;
    q.setIdentity();
    g2o::SE3Quat pose(q, trans);
    g2o::VertexSE3Expmap *v_se3 = new g2o::VertexSE3Expmap();
    v_se3->setId(vertex_id++);
    v_se3->setEstimate(pose);
    optimizer.addVertex(v_se3);

    // set camera intrinsic
    cv::Mat _K;
    K.convertTo(_K, CV_64FC1);
    g2o::CameraParameters *cam_params = new g2o::CameraParameters(_K.at<double>(0, 0),
                                                                  Eigen::Vector2d(_K.at<double>(0, 2),
                                                                                  _K.at<double>(1, 2)), 0.);
    cam_params->setId(0);
    optimizer.addParameter(cam_params);

    // add vertex for points and edges for projection
    for (int i = 0; i < 0; i++) {
        const auto &pt = pts3d[i];
        const auto &uv = pts2d[i];
        // printf("%.3f %.3f %.3f %.3f %.3f\n",pt.x, pt.y, pt.z, uv.x,uv.y);

        g2o::VertexSBAPointXYZ *v_p = new g2o::VertexSBAPointXYZ();
        v_p->setId(vertex_id++);
        v_p->setMarginalized(true);
        v_p->setEstimate(Eigen::Vector3d(pt.x, pt.y, pt.z));
        v_p->setFixed(true);
        optimizer.addVertex(v_p);

        g2o::EdgeProjectXYZ2UV *e = new g2o::EdgeProjectXYZ2UV();
        e->setVertex(0, v_p);
        e->setVertex(1, v_se3);
        e->setMeasurement(Eigen::Vector2d(uv.x, uv.y));
        e->information() = Eigen::Matrix2d::Identity();
        e->setRobustKernel(new g2o::RobustKernelHuber);
        e->setParameterId(0, 0);
        optimizer.addEdge(e);
    }

    // add vertex for lines and edges for projection
    for (int i = 0; i < nlns; i++) {
        const auto &ln3d = lns3d[i];
        const auto &ln2d = lns2d[i];
        // printf("%.3f %.3f %.3f %.3f %.3f\n",pt.x, pt.y, pt.z, uv.x,uv.y);

        g2o::VertexPluckerLine *v_l = new g2o::VertexPluckerLine();
        v_l->setId(vertex_id++);
        v_l->setMarginalized(true);
        g2o::Vector6D temp;
        g2o::Vector3D start_temp;
        g2o::Vector3D end_temp;

        start_temp << ln3d[0], ln3d[1], ln3d[2];
        end_temp << ln3d[3], ln3d[4], ln3d[5];
        g2o::Vector3D tail3 = end_temp - start_temp;
        g2o::Vector3D head3 = start_temp.cross(end_temp);

        temp.head<3>() = head3;
        temp.tail<3>() = tail3;
        v_l->setEstimate(temp);
        v_l->setFixed(true);
        optimizer.addVertex(v_l);

        g2o::EdgePluckerLine *e = new g2o::EdgePluckerLine();
        e->setVertex(0, v_l);
        e->setVertex(1, v_se3);
        e->setMeasurement(Eigen::Vector4d(ln2d[0], ln2d[1], ln2d[2], ln2d[3]));
        e->information() = Eigen::Matrix4d::Identity();
        e->setRobustKernel(new g2o::RobustKernelHuber);
        e->setParameterId(0, 0);
        optimizer.addEdge(e);
    }

    // optimize
    optimizer.initializeOptimization();
    optimizer.optimize(10);
    // output result
    Eigen::MatrixXd T = Eigen::Isometry3d(v_se3->estimate()).matrix();
    R = (cv::Mat_<float>(3, 3) << T(0, 0), T(0, 1), T(0, 2),
            T(1, 0), T(1, 1), T(1, 2),
            T(2, 0), T(2, 1), T(2, 2));
    t = (cv::Mat_<float>(3, 1) << T(0, 3), T(1, 3), T(2, 3));
}

void Optimize::PnPL(const std::vector<cv::Point3f> &pts3d, const std::vector<cv::Point2f> &pts2d,
                    const std::vector<cv::Vec6f> &lns3d, const std::vector<cv::Vec4f> &lns2d,
                    const cv::Mat &K, cv::Mat &R, cv::Mat &t, bool line_constraint) {
    int npts = pts3d.size();
    int nlns = lns3d.size();

    // init g2o
    g2o::BlockSolverX::LinearSolverType *linearSolver;
    linearSolver = new g2o::LinearSolverCSparse<g2o::BlockSolverX::PoseMatrixType>();
    g2o::BlockSolverX *solver_ptr = new g2o::BlockSolverX(linearSolver);
    g2o::OptimizationAlgorithmLevenberg *solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    g2o::SparseOptimizer optimizer;
    // optimizer.setVerbose(true);
    optimizer.setAlgorithm(solver);

    // add vertex for pose
    int vertex_id = 0;
    Eigen::Vector3d trans(0.1, 0.1, 0.1);
    Eigen::Quaterniond q;
    q.setIdentity();
    g2o::SE3Quat pose(q, trans);
    g2o::VertexSE3Expmap *v_se3 = new g2o::VertexSE3Expmap();
    v_se3->setId(vertex_id++);
    v_se3->setEstimate(pose);
    optimizer.addVertex(v_se3);

    // set camera intrinsic
    cv::Mat _K;
    K.convertTo(_K, CV_64FC1);
    g2o::CameraParameters *cam_params = new g2o::CameraParameters(_K.at<double>(0, 0),
                                                                  Eigen::Vector2d(_K.at<double>(0, 2),
                                                                                  _K.at<double>(1, 2)), 0.);
    cam_params->setId(0);
    optimizer.addParameter(cam_params);

    // add vertex for points and edges for projection

    for (int i = 0; i < npts; i++) {
        const auto &pt = pts3d[i];
        const auto &uv = pts2d[i];
        // printf("%.3f %.3f %.3f %.3f %.3f\n",pt.x, pt.y, pt.z, uv.x,uv.y);

        g2o::VertexSBAPointXYZ *v_p = new g2o::VertexSBAPointXYZ();
        v_p->setId(vertex_id++);
        v_p->setMarginalized(true);
        v_p->setEstimate(Eigen::Vector3d(pt.x, pt.y, pt.z));
        v_p->setFixed(true);
        optimizer.addVertex(v_p);

        g2o::EdgeProjectXYZ2UV *e = new g2o::EdgeProjectXYZ2UV();
        e->setVertex(0, v_p);
        e->setVertex(1, v_se3);
        e->setMeasurement(Eigen::Vector2d(uv.x, uv.y));
        e->information() = Eigen::Matrix2d::Identity();
        e->setRobustKernel(new g2o::RobustKernelHuber);
        e->setParameterId(0, 0);
        optimizer.addEdge(e);
    }

    if (line_constraint) {

        // add vertex for lines and edges for projection
        for (int i = 0; i < nlns; i++) {
            const auto &ln3d = lns3d[i];
            const auto &ln2d = lns2d[i];
            // printf("%.3f %.3f %.3f %.3f %.3f\n",pt.x, pt.y, pt.z, uv.x,uv.y);

            g2o::VertexSBALine *v_l = new g2o::VertexSBALine();
            v_l->setId(vertex_id++);
            v_l->setMarginalized(true);
            g2o::Vector6D temp;
            temp << ln3d[0], ln3d[1], ln3d[2], ln3d[3], ln3d[4], ln3d[5];
            v_l->setEstimate(temp);
            v_l->setFixed(true);
            optimizer.addVertex(v_l);

            g2o::EdgeProjectLine *e = new g2o::EdgeProjectLine();
            e->setVertex(0, v_l);
            e->setVertex(1, v_se3);
            e->setMeasurement(Eigen::Vector4d(ln2d[0], ln2d[1], ln2d[2], ln2d[3]));
            e->information() = Eigen::Matrix4d::Identity();
            e->setRobustKernel(new g2o::RobustKernelHuber);
            e->setParameterId(0, 0);
            optimizer.addEdge(e);
        }

        for (int i = 0; i < nlns / 2; i++) {
            const auto &ln3d = lns3d[i];
            const auto &ln2d = lns2d[i];

            const auto &ln3d_pt = lns3d[nlns / 2 + i];
            // printf("%.3f %.3f %.3f %.3f %.3f\n",pt.x, pt.y, pt.z, uv.x,uv.y);

            g2o::VertexSBALine *v_l = new g2o::VertexSBALine();
            v_l->setId(vertex_id++);
//            v_l->setMarginalized(true);
            g2o::Vector6D temp;
            temp << ln3d[0], ln3d[1], ln3d[2], ln3d[3], ln3d[4], ln3d[5];
            v_l->setEstimate(temp);
            v_l->setFixed(true);
            optimizer.addVertex(v_l);

            g2o::VertexSBAPointXYZ *v_p1 = new g2o::VertexSBAPointXYZ();
            v_p1->setId(vertex_id++);
//            v_p1->setMarginalized(true);
            v_p1->setEstimate(Eigen::Vector3d(ln3d_pt[0], ln3d_pt[1], ln3d_pt[2]));
//            v_p1->setFixed(true);
            optimizer.addVertex(v_p1);

            g2o::VertexSBAPointXYZ *v_p2 = new g2o::VertexSBAPointXYZ();
            v_p2->setId(vertex_id++);
//            v_p2->setMarginalized(true);
            v_p2->setEstimate(Eigen::Vector3d(ln3d_pt[3], ln3d_pt[4], ln3d_pt[5]));
//            v_p2->setFixed(true);
            optimizer.addVertex(v_p2);

            g2o::EdgeLineSBAPoint *e1 = new g2o::EdgeLineSBAPoint();
            e1->setVertex(0, v_l);
            e1->setVertex(1, v_p1);
            e1->setMeasurement(Eigen::Vector2d(0, 0));
            e1->information() = Eigen::Matrix2d::Identity();
            e1->setRobustKernel(new g2o::RobustKernelHuber);
            e1->setParameterId(0, 0);
            optimizer.addEdge(e1);

            g2o::EdgeLineSBAPoint *e2 = new g2o::EdgeLineSBAPoint();
            e2->setVertex(0, v_l);
            e2->setVertex(1, v_p2);
            e2->setMeasurement(Eigen::Vector2d(0, 0));
            e2->information() = Eigen::Matrix2d::Identity();
            e2->setRobustKernel(new g2o::RobustKernelHuber);
            e2->setParameterId(0, 0);
            optimizer.addEdge(e2);
        }
    }


    // optimize
    optimizer.initializeOptimization();
    optimizer.optimize(10);
    // output result
    Eigen::MatrixXd T = Eigen::Isometry3d(v_se3->estimate()).matrix();
    R = (cv::Mat_<float>(3, 3) << T(0, 0), T(0, 1), T(0, 2),
            T(1, 0), T(1, 1), T(1, 2),
            T(2, 0), T(2, 1), T(2, 2));
    t = (cv::Mat_<float>(3, 1) << T(0, 3), T(1, 3), T(2, 3));
}


void Optimize::bundleAdjust(Eigen::Matrix3d &R, Eigen::Vector3d &t, bool useStruct) {
    int npts = m3dPoints.size();

    // init g2o
    g2o::BlockSolverX::LinearSolverType *linearSolver;
    linearSolver = new g2o::LinearSolverCSparse<g2o::BlockSolverX::PoseMatrixType>();
    g2o::BlockSolverX *solver_ptr = new g2o::BlockSolverX(linearSolver);
    g2o::OptimizationAlgorithmLevenberg *solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    g2o::SparseOptimizer optimizer;
//    optimizer.setVerbose(true);
    optimizer.setAlgorithm(solver);

    // add vertex for pose
    int vertex_id = 0;
    Eigen::Vector3d trans(0.1, 0.1, 0.1);
    Eigen::Quaterniond q;
    q.setIdentity();
    g2o::SE3Quat pose(q, trans);
    g2o::VertexSE3Expmap *v_se3 = new g2o::VertexSE3Expmap();
    v_se3->setId(vertex_id++);
    v_se3->setEstimate(pose);
    optimizer.addVertex(v_se3);

    // set camera intrinsic
    double f;
    Eigen::Vector2d cxy;
    f = K(0, 0);
    cxy = K.block<2, 1>(0, 2);
//    std::cout << cxy << std::endl;
    g2o::CameraParameters *cam_params = new g2o::CameraParameters(f, cxy, 0.);
    cam_params->setId(0);

    const int line_start_id = vertex_id;
    optimizer.addParameter(cam_params);
    if (useStruct) {
        for (int j = 0; j < mLines.size(); ++j) {
            g2o::VertexSBALine *v_line = new g2o::VertexSBALine();
            v_line->setId(vertex_id++);
            v_line->setMarginalized(false);
            v_line->setEstimate(mLines[j]);
            optimizer.addVertex(v_line);
        }
    }

    const int plane_start_id = vertex_id;
    if (useStruct) {
        for (int k = 0; k < mPlanes.size(); ++k) {
            g2o::VertexSBAPlane *v_plane = new g2o::VertexSBAPlane();
            v_plane->setId(vertex_id++);
            v_plane->setMarginalized(false);
            v_plane->setEstimate(mPlanes[k]);
            optimizer.addVertex(v_plane);
        }
    }

    // add vertex for points and edges for projection
    const int point_start_d = vertex_id;
    for (int i = 0; i < npts; i++) {
        const auto &pt = m3dPoints[i];
        const auto &uv = m2dPoints[i];

        g2o::VertexSBAPointXYZ *v_p = new g2o::VertexSBAPointXYZ();
        v_p->setId(vertex_id++);
        v_p->setMarginalized(false);
        v_p->setEstimate(pt);
        optimizer.addVertex(v_p);

        g2o::EdgeProjectXYZ2UV *e = new g2o::EdgeProjectXYZ2UV();
        e->setVertex(0, v_p);
        e->setVertex(1, v_se3);
        e->setMeasurement(uv);
        e->information() = Eigen::Matrix2d::Identity();
        e->setRobustKernel(new g2o::RobustKernelHuber);
        e->setParameterId(0, 0);
        optimizer.addEdge(e);
        if (useStruct) {
            if (mPointCate[i] == LinePoint) {
                g2o::EdgeLineSBAPoint *e_line = new g2o::EdgeLineSBAPoint();
                e_line->setVertex(1, v_p);
                e_line->setVertex(0, optimizer.vertex(line_start_id + mPointLineIndex[i]));
                e_line->setMeasurement(Eigen::Vector2d(0, 0));
                e_line->information() = Eigen::Matrix2d::Identity();
                e_line->setRobustKernel(new g2o::RobustKernelHuber);
                e_line->setParameterId(0, 0);
                optimizer.addEdge(e_line);
            }

            if (mPointCate[i] == PlanePoint) {
                g2o::EdgePlaneSBAPoint *e_plane = new g2o::EdgePlaneSBAPoint();
                e_plane->setVertex(1, v_p);
                e_plane->setVertex(0, optimizer.vertex(plane_start_id + mPointPlaneIndex[i]));
                e_plane->setMeasurement(0);

                e_plane->setRobustKernel(new g2o::RobustKernelHuber);
                e_plane->setParameterId(0, 0);
                optimizer.addEdge(e_plane);

            }
        }
    }

    // optimize
    optimizer.initializeOptimization();
    optimizer.optimize(20);
    // output result
    Eigen::MatrixXd T = Eigen::Isometry3d(v_se3->estimate()).matrix();
    R = T.block<3, 3>(0, 0);
    t = T.block<3, 1>(0, 3);
}

